#!/bin/bash

# ================================ 配置区域 ================================
# 模型路径配置
MODEL_CKPT="/root/bwh/Models/mistralai_Mistral-7B-v0.1"

# 实验配置
DATASET_NAME="MATH"
START_IDX=0
END_IDX=50
NUM_ROLLOUTS=32
NOTE="mistral_7b_math_verbose"

# GPU配置
GPU_IDS="0,1,2,3,4,5,6,7"
TENSOR_PARALLEL_SIZE=8

# 输出目录配置
OUTPUT_ROOT="./run_outputs"
TIMESTAMP=$(date +"%Y-%m-%d_%H-%M-%S")
EXP_DIR="${OUTPUT_ROOT}/${DATASET_NAME}/$(basename ${MODEL_CKPT})/${TIMESTAMP}---${NOTE}"

# ================================ 环境准备 ================================
echo "=========================================="
echo "rStar数学推理生成器 - MATH数据集版本"
echo "=========================================="
echo "时间: $(date)"
echo "实验目录: ${EXP_DIR}"
echo "模型: ${MODEL_CKPT}"
echo "数据集: ${DATASET_NAME}"
echo "问题范围: ${START_IDX}-${END_IDX}"
echo "MCTS轮数: ${NUM_ROLLOUTS}"
echo "最大深度: 8"
echo "GPU配置: ${GPU_IDS} (张量并行: ${TENSOR_PARALLEL_SIZE})"
echo "=========================================="

# 第一步：彻底重置GPU状态
echo "正在重置GPU状态..."
python3 scripts/reset_gpu.py

# 验证重置结果
if [ $? -ne 0 ]; then
    echo "❌ GPU重置失败"
    exit 1
fi
echo "✅ GPU重置完成"

# 第二步：设置vLLM环境变量
echo "正在设置vLLM环境变量..."
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export VLLM_DISABLE_STRICT_MEMORY_CHECK=1
export VLLM_GPU_MEMORY_UTILIZATION=0.85
export VLLM_ALLOW_RUNTIME_LORA_UPDATING=1
export VLLM_TORCH_COMPILE_LEVEL=0
export VLLM_DISABLE_CUSTOM_ALL_REDUCE=1
export CUDA_LAUNCH_BLOCKING=0
echo "✅ vLLM环境变量设置完成"

# 第三步：最终环境检查
echo "正在进行最终环境检查..."
python3 -c "
import sys
sys.path.append('.')
from common.gpu_utils import check_memory_requirements

# 检查内存要求（每个GPU至少需要8GB可用内存用于Mistral-7B）
print('检查内存要求...')
if not check_memory_requirements(8.0, [0,1,2,3,4,5,6,7]):
    print('❌ 内存不足，无法运行模型')
    exit(1)

print('✅ 环境检查完成')
"

# 验证环境检查结果
if [ $? -ne 0 ]; then
    echo "❌ 环境检查失败"
    exit 1
fi

# 第四步：创建输出目录
echo "正在创建输出目录..."
mkdir -p "${EXP_DIR}"
mkdir -p "${EXP_DIR}/answer_sheets"
mkdir -p "${EXP_DIR}/logs"
echo "✅ 输出目录创建完成: ${EXP_DIR}"

# 第五步：保存实验配置
echo "正在保存实验配置..."
cat > "${EXP_DIR}/experiment_config.txt" << EOF
实验配置信息
=====================================
开始时间: $(date)
模型路径: ${MODEL_CKPT}
数据集: ${DATASET_NAME}
问题范围: ${START_IDX} - ${END_IDX}
MCTS轮数: ${NUM_ROLLOUTS}
最大深度: 8
GPU配置: ${GPU_IDS}
张量并行: ${TENSOR_PARALLEL_SIZE}
实验标签: ${NOTE}
输出目录: ${EXP_DIR}
=====================================
EOF

# ================================ 主程序执行 ================================
echo "🚀 开始运行生成任务..."

# 运行主程序，添加所有必要的输出和调试参数
CUDA_VISIBLE_DEVICES=${GPU_IDS} python3 run_src/do_generate.py \
    --model_ckpt "${MODEL_CKPT}" \
    --api vllm \
    --tensor_parallel_size ${TENSOR_PARALLEL_SIZE} \
    --max_num_seqs 32 \
    --max_tokens 1024 \
    --dataset_name ${DATASET_NAME} \
    --data_root "./data" \
    --test_json_filename "test_all" \
    --prompts_root "./prompts" \
    --outputs_dir "${OUTPUT_ROOT}" \
    --num_rollouts ${NUM_ROLLOUTS} \
    --mcts_exploration_weight 2.0 \
    --mcts_weight_scheduler const \
    --mcts_discount_factor 1.0 \
    --num_subquestions 3 \
    --num_votes 10 \
    --max_depth_allowed 8 \
    --start_idx ${START_IDX} \
    --end_idx ${END_IDX} \
    --verbose \
    --save_tree \
    --note "${NOTE}" \
    --seed 42 \
    2>&1 | tee "${EXP_DIR}/logs/generation.log"

# 检查执行结果
EXIT_CODE=$?
if [ $EXIT_CODE -eq 0 ]; then
    echo "✅ 生成任务执行成功！"
else
    echo "❌ 生成任务执行失败，退出码: $EXIT_CODE"
fi

# ================================ 结果统计 ================================
echo "正在统计结果..."

# 统计生成的文件数量
ANSWER_FILES=$(find "${EXP_DIR}/answer_sheets" -name "*Answer.json" | wc -l)
SOLUTION_FILES=$(find "${EXP_DIR}/answer_sheets" -name "*Solutions.json" | wc -l)
TREE_FILES=$(find "${EXP_DIR}/answer_sheets" -name "*.tree" | wc -l)

echo "=========================================="
echo "执行结果统计"
echo "=========================================="
echo "Answer文件数量: ${ANSWER_FILES}"
echo "Solution文件数量: ${SOLUTION_FILES}"
echo "Tree文件数量: ${TREE_FILES}"
echo "日志文件: ${EXP_DIR}/logs/generation.log"
echo "实验目录: ${EXP_DIR}"
echo "完成时间: $(date)"
echo "=========================================="

# 保存最终统计信息
cat > "${EXP_DIR}/final_summary.txt" << EOF
实验完成统计
=====================================
完成时间: $(date)
退出码: ${EXIT_CODE}
生成文件统计:
  - Answer文件: ${ANSWER_FILES}
  - Solution文件: ${SOLUTION_FILES}
  - Tree文件: ${TREE_FILES}
实验目录: ${EXP_DIR}
=====================================
EOF

# 显示输出目录结构
echo "输出目录结构:"
ls -la "${EXP_DIR}/answer_sheets" | head -20

# 如果有tree文件，显示一个示例
if [ ${TREE_FILES} -gt 0 ]; then
    echo "Tree文件示例:"
    FIRST_TREE=$(find "${EXP_DIR}/answer_sheets" -name "*.tree" | head -1)
    if [ -n "${FIRST_TREE}" ]; then
        echo "文件: $(basename ${FIRST_TREE})"
        head -20 "${FIRST_TREE}"
    fi
fi

echo "🎉 任务执行完成！"
echo "📁 查看结果: ${EXP_DIR}"
echo "📊 查看日志: ${EXP_DIR}/logs/generation.log" 